
import numpy as np
from model_definition import RNN, LSTM
import torch


def lstm_predict(x):
    # 输入的维度为1，只有Close收盘价
    input_dim = 1
    # 隐藏层特征的维度
    hidden_dim = 32
    # 循环的layers
    num_layers = 2
    # 预测后一天的收盘价
    output_dim = 1

    lstm = LSTM(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim, num_layers=num_layers)
    lstm.load_state_dict(torch.load('model/lstm_model.bin'))

    x_hat = np.mean(x)
    x_std = np.std(x)
    input_data = (x - x_hat) / x_std
    output = lstm(torch.unsqueeze(torch.Tensor(input_data), dim=0).view(-1, len(input_data), 1))
    # output = lstm(torch.Tensor(input_data))

    output_value = torch.squeeze(output).cpu().detach().numpy()
    prediction = x_hat + output_value
    return prediction


if __name__ == '__main__':
    x = np.array([1650.3, 1646.3, 1633, 1654.2, 1620.3, 1616.3, 1611.1, 1622.2, 1646.3, 1633])
    # input_data = torch.unsqueeze(torch.Tensor(input_data), dim=0).view(-1, len(input_data), 1)
    # print(input_data.shape)
    # x = np.array([[[1650.3], [1646.3], [1633], [1654.2], [1620.3],[ 1616.3], [1611.1], [1622.2], [1646.3], [1633]]])
    # print(x.shape)
    lstm_predict(x, 1, 1, 1)
